Track parquet writer encoding memory usage on MemoryPool #11345

wiedld · 2024-07-09T02:33:16Z

Which issue does this PR close?

Rationale for this change

ParquetSink can use non-trivial amounts of memory to buffer rowgroups prior to flush, when executed within a task context. Therefore, this memory usage should be accounted within the task's memory pool.

What changes are included in this PR?

Ensure memory accounting under three use cases for ParquetSink:

when non-parallel writes
when using multiple row-writers
when using multiple column-writers

How parallelized-write memory tracking works, is summarized here. I feel like it should be a doc comment, but I wasn't sure where. 🤔

Are these changes tested?

Yes

Are there any user-facing changes?

No

datafusion/core/src/datasource/file_format/parquet.rs

…n, before selecting shrinking for data bytes flushed

alamb

Thank you @wiedld

For anyone else reading this PR, we discussed this face to face and I think @wiedld has some ideas of how to make it simpler

datafusion/core/src/datasource/file_format/parquet.rs

…arallelized call stack for col vs rg

alamb · 2024-07-09T22:02:04Z

Thanks @wiedld -- please mark this PR as ready for review when it is ready for another look

wiedld · 2024-07-10T00:20:21Z

datafusion/core/tests/memory_limit/mod.rs

+            // TODO: update error handling in ParquetSink
+            "Unable to send array to writer!",


The parallelized writes have vectors for channels, and vectors for the spawned tasks. This error we are hitting (on memory limit reached) is for the closed channel.

I believe we want to be surfacing errors for the tasks, which should exit due to the memory reservation error. Need to poke around a bit more.

I think we need to update several map_err statements to propagate inner error messages rather than ignore them. E.g.

datafusion/datafusion/core/src/datasource/file_format/parquet.rs

Lines 880 to 884 in b96186f

col_array_channels[next_channel]

.send(c)

.await

.map_err(|_| {

DataFusionError::Internal("Unable to send array to writer!".into())

change to something like

col_array_channels[next_channel] .send(c) .await .map_err(|e| internal_datafusion_err!("Unable to send array to writer due to error {e}"))

Filed #11397 to track

alamb · 2024-07-10T11:22:05Z

FYI @devinjdangelo

alamb · 2024-07-10T11:23:19Z

datafusion/core/src/datasource/file_format/parquet.rs

+            let (writer, col_reservation) = task.join_unwind().await?;
+            let encoded_size = writer.get_estimated_total_bytes();
+            rg_reservation.grow(encoded_size);
+            drop(col_reservation);


I don't think the explicit drop is needed -- it will be dropped automatically by the compiler when col_reservation goes out of scope on the next line

alamb

Thanks @wiedld -- this is looking quite close

datafusion/core/src/datasource/file_format/parquet.rs

devinjdangelo

Thanks @wiedld and @alamb this looks good to me! I left a comment on how I think we can improve the error messages.

devinjdangelo · 2024-07-10T16:07:58Z

datafusion/core/tests/memory_limit/mod.rs

+            // TODO: update error handling in ParquetSink
+            "Unable to send array to writer!",


I think we need to update several map_err statements to propagate inner error messages rather than ignore them. E.g.

datafusion/datafusion/core/src/datasource/file_format/parquet.rs

Lines 880 to 884 in b96186f

col_array_channels[next_channel]

.send(c)

.await

.map_err(|_| {

DataFusionError::Internal("Unable to send array to writer!".into())

change to something like

col_array_channels[next_channel] .send(c) .await .map_err(|e| internal_datafusion_err!("Unable to send array to writer due to error {e}"))

wiedld · 2024-07-10T16:32:10Z

Thanks @wiedld and @alamb this looks good to me! I left a comment on how I think we can improve the error messages.

Thank you @devinjdangelo ! Note that the error message propagated back is channel closed and not the resource exhuastion. I've been fiddling a bit with a solution, and I'll put up a PR (pinging you both) ~~hopefully later today~~ right after this merges.

alamb

Thank you @wiedld and @devinjdangelo -- I think this looks good to me now

I agree with @devinjdangelo it would be great to improve the error message propagation -- but I also think we could do that as a follow on PR.

Edit: looks like @wiedld plans to do as a follow on PR so I'll file a ticket and merge this one

Follow on ticket: #11397

alamb · 2024-07-10T18:21:09Z

🚀

* feat(11344): track memory used for non-parallel writes * feat(11344): track memory usage during parallel writes * test(11344): create bounded stream for testing * test(11344): test ParquetSink memory reservation * feat(11344): track bytes in file writer * refactor(11344): tweak the ordering to add col bytes to rg_reservation, before selecting shrinking for data bytes flushed * refactor: move each col_reservation and rg_reservation to match the parallelized call stack for col vs rg * test(11344): add memory_limit enforcement test for parquet sink * chore: cleanup to remove unnecessary reservation management steps * fix: fix CI test failure due to file extension rename

feat(11344): track memory used for non-parallel writes

4eeca45

github-actions bot added the core Core DataFusion crate label Jul 9, 2024

wiedld commented Jul 9, 2024

View reviewed changes

datafusion/core/src/datasource/file_format/parquet.rs Show resolved Hide resolved

wiedld added 3 commits July 8, 2024 19:54

feat(11344): track memory usage during parallel writes

07e45e4

test(11344): create bounded stream for testing

68beda4

test(11344): test ParquetSink memory reservation

073e471

wiedld force-pushed the 11344/track-parquet-sink-memory branch from 518a7e9 to 073e471 Compare July 9, 2024 02:54

wiedld commented Jul 9, 2024

View reviewed changes

datafusion/core/src/datasource/file_format/parquet.rs Outdated Show resolved Hide resolved

wiedld commented Jul 9, 2024

View reviewed changes

datafusion/core/src/datasource/file_format/parquet.rs Show resolved Hide resolved

feat(11344): track bytes in file writer

3709d48

wiedld commented Jul 9, 2024

View reviewed changes

datafusion/core/src/datasource/file_format/parquet.rs Show resolved Hide resolved

refactor(11344): tweak the ordering to add col bytes to rg_reservatio…

f880fae

…n, before selecting shrinking for data bytes flushed

wiedld marked this pull request as ready for review July 9, 2024 03:59

alamb reviewed Jul 9, 2024

View reviewed changes

alamb marked this pull request as draft July 9, 2024 19:27

refactor: move each col_reservation and rg_reservation to match the p…

558348e

…arallelized call stack for col vs rg

test(11344): add memory_limit enforcement test for parquet sink

eaa5925

wiedld commented Jul 10, 2024

View reviewed changes

alamb reviewed Jul 10, 2024

View reviewed changes

datafusion/core/src/datasource/file_format/parquet.rs Show resolved Hide resolved

datafusion/core/src/datasource/file_format/parquet.rs Outdated Show resolved Hide resolved

chore: cleanup to remove unnecessary reservation management steps

6ce964e

wiedld force-pushed the 11344/track-parquet-sink-memory branch from 46dc2d3 to 6ce964e Compare July 10, 2024 15:50

devinjdangelo reviewed Jul 10, 2024

View reviewed changes

fix: fix CI test failure due to file extension rename

fcf1476

wiedld marked this pull request as ready for review July 10, 2024 16:58

alamb approved these changes Jul 10, 2024

View reviewed changes

alamb mentioned this pull request Jul 10, 2024

Improve error messages for parallel parquet writer "Unable to send array to writer!" #11397

Closed

alamb changed the title ~~feat(11344): track parquet encoding memory usage~~ Track parquet writer encoding memory usage on MemoryPool Jul 10, 2024

alamb merged commit 6038f4c into apache:main Jul 10, 2024
24 checks passed

alamb deleted the 11344/track-parquet-sink-memory branch July 10, 2024 18:21

appletreeisyellow mentioned this pull request Jul 12, 2024

WIP(iox-11398): patched df upgrade 2024-07-08 influxdata/arrow-datafusion#33

Closed

appletreeisyellow mentioned this pull request Jul 22, 2024

WIP(iox-11398): patched df upgrade 2024-07-13 influxdata/arrow-datafusion#35

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Track parquet writer encoding memory usage on MemoryPool #11345

Track parquet writer encoding memory usage on MemoryPool #11345

wiedld commented Jul 9, 2024 •

edited

Loading

alamb left a comment

alamb commented Jul 9, 2024

wiedld Jul 10, 2024

devinjdangelo Jul 10, 2024

alamb Jul 10, 2024

alamb commented Jul 10, 2024

alamb Jul 10, 2024

alamb left a comment

devinjdangelo left a comment

devinjdangelo Jul 10, 2024

wiedld commented Jul 10, 2024 •

edited

Loading

alamb left a comment •

edited

Loading

alamb commented Jul 10, 2024

		// TODO: update error handling in ParquetSink
		"Unable to send array to writer!",

	col_array_channels[next_channel]
	.send(c)
	.await
	.map_err(\|_\| {
	DataFusionError::Internal("Unable to send array to writer!".into())

Track parquet writer encoding memory usage on MemoryPool #11345

Track parquet writer encoding memory usage on MemoryPool #11345

Conversation

wiedld commented Jul 9, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

alamb left a comment

Choose a reason for hiding this comment

alamb commented Jul 9, 2024

wiedld Jul 10, 2024

Choose a reason for hiding this comment

devinjdangelo Jul 10, 2024

Choose a reason for hiding this comment

alamb Jul 10, 2024

Choose a reason for hiding this comment

alamb commented Jul 10, 2024

alamb Jul 10, 2024

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

devinjdangelo left a comment

Choose a reason for hiding this comment

devinjdangelo Jul 10, 2024

Choose a reason for hiding this comment

wiedld commented Jul 10, 2024 • edited Loading

alamb left a comment • edited Loading

Choose a reason for hiding this comment

alamb commented Jul 10, 2024

wiedld commented Jul 9, 2024 •

edited

Loading

wiedld commented Jul 10, 2024 •

edited

Loading

alamb left a comment •

edited

Loading